#include "KMultiplexor.h"
#include "KKey.h"
#include "KPackage.h"
#include "KSockContext.h"
#include "KSocketIocp.h"
#include "KStorage.h"
#include "KTaskThread.h"
#include <algorithm>

KMultiplexor::KMultiplexor(size_t numOfWorkers) :
    m_stopped(true),
    m_gc(new KTaskThread),
    m_workers(new KTaskThread[numOfWorkers]),
    m_pool(KDef::DefaultPoolSize),
    m_chunkSize(KDef::DefaultChunkSize),
    m_nThreads(numOfWorkers)
{
    assert(numOfWorkers > 0);
    KTaskList::TaskType task(std::bind(&KMultiplexor::_worker, this));
    while (numOfWorkers--) {
        m_workers[numOfWorkers].append(task);
    }
    m_gc->append(std::bind(&KMultiplexor::_gc, this));
}

KMultiplexor::~KMultiplexor()
{
    stop();
    assert(m_stopped);
}

KSocket *KMultiplexor::addAcceptor(const KIAddress &addrLocal, KKey *key, int sendBufSize, bool reuseAddrLocal, size_t extra)
{
    if (m_stopped) {
        onError(NULL, KDef::ErrorNotStarted, 0);
        delete key;
        return NULL;
    }
    KAutoPtr<KTcpSocketContext> acceptor(new KTcpSocketContext);
    acceptor->setKey(key);
    if (acceptor->open(sendBufSize) &&
        (reuseAddrLocal ? acceptor->setReuseAddr(true) : true) &&
        acceptor->bind(addrLocal) &&
        acceptor->listen(SOMAXCONN) &&
        acceptor->initAcceptEx()
    ) {
        if (!m_iocp->addSocket(acceptor, reinterpret_cast<ULONG_PTR>(static_cast<KDef::SocketContext *>(acceptor)))) {
            DWORD lastErr = GetLastError();
            KDebugLib("Failed to add acceptor to i/o completion port: %d", lastErr);
            onError(acceptor, KDef::ErrorAddToIocp, lastErr);
            return NULL;
        }
        size_t n = 0;
        m_context.pushFront(acceptor);
        for (; n <= extra && accept(acceptor); ++n);
        if (n > 0) {
            return acceptor.release();
        }
        m_context.cutOff(acceptor);
        return NULL;
    }
    KDebugLib("Acceptor failed to initialize");
    onError(acceptor, KDef::ErrorAcceptorInit, WSAGetLastError());
    return NULL;
}

KSocket *KMultiplexor::addConnector(const KIAddress &addrLocal, const KIAddress &addrRemote, KKey *key, int sendBufSize, bool reuseAddrLocal)
{
    if (m_stopped) {
        onError(NULL, KDef::ErrorNotStarted, 0);
        delete key;
        return NULL;
    }
    KAutoPtr<KTcpSocketContext> connector(new KTcpSocketContext);
    connector->setKey(key);
    if (connector->open(sendBufSize) &&
        (reuseAddrLocal ? connector->setReuseAddr(true) : true) &&
        connector->bind(addrLocal) &&
        connector->initConnectEx() &&
        connector->initDisconnectEx()
    ) {
        if (!m_iocp->addSocket(connector, reinterpret_cast<ULONG_PTR>(static_cast<KDef::SocketContext *>(connector)))) {
            DWORD lastErr = GetLastError();
            KDebugLib("Failed to add connector to i/o completion port: %d", lastErr);
            onError(connector, KDef::ErrorAddToIocp, lastErr);
            return NULL;
        }
        m_context.pushFront(connector);
        if (connect(connector, addrRemote)) {
            return connector.release();
        }
        m_context.cutOff(connector);
        return NULL;
    }
    KDebugLib("Connector failed to initialize");
    onError(connector, KDef::ErrorConnectorInit, WSAGetLastError());
    return NULL;
}

KSocket *KMultiplexor::addMessenger(const KIAddress &addrLocal, KKey *key, int sendBufSize, bool reuseAddrLocal)
{
    if (m_stopped) {
        onError(NULL, KDef::ErrorNotStarted, 0);
        delete key;
        return NULL;
    }
    KAutoPtr<KUdpSocketContext> messenger(new KUdpSocketContext);
    messenger->setKey(key);
    if (messenger->open(sendBufSize) &&
        (reuseAddrLocal ? messenger->setReuseAddr(true) : true) &&
        messenger->bind(addrLocal)
    ) {
        if (!m_iocp->addSocket(messenger, reinterpret_cast<ULONG_PTR>(static_cast<KDef::SocketContext *>(messenger)))) {
            DWORD lastErr = GetLastError();
            KDebugLib("Failed to add messenger to i/o completion port: %d", lastErr);
            onError(messenger, KDef::ErrorAddToIocp, lastErr);
            return NULL;
        }
        m_context.pushFront(messenger);
        onCreated(messenger);
        return messenger.release();
    }
    KDebugLib("Messenger failed to initialize");
    onError(messenger, KDef::ErrorMessengerInit, WSAGetLastError());
    return NULL;
}

bool KMultiplexor::cancel(KSocket *sock, KDef::BasicRequest *request)
{
    deleteLater(sock);
    return sock->asyncCancel(request);
}

bool KMultiplexor::postMessage(KSocket *sock, KKey *key)
{
    assert(sock != NULL);
    KAutoPtr<KDef::MessageRequest> request(new KDef::MessageRequest(key));
    KDef::SocketContext *node = sock->getContext();
    node->addRequest(request);
    if (m_iocp->postStatus(KSocketIocp::Status(0, reinterpret_cast<ULONG_PTR>(node), request))) {
        request.release();
        node->resetIdleCount();
        return true;
    }
    onError(sock, KDef::ErrorPostStatus, GetLastError());
    node->removeRequest(request);
    return false;
}

bool KMultiplexor::receive(KSocket *sock, size_t length, KKey *key)
{
    assert(sock != NULL);
    DWORD lastErr;
    KDef::IoRequest *request = static_cast<KDef::IoRequest *>(m_pool.popFront());
    KDef::SocketContext *node = sock->getContext();
    length = std::max<size_t>(length, m_chunkSize);
    if (request) {
        request->reset();
        request->resize(length);
    } else {
        request = new KDef::IoRequest(length);
        assert(request != NULL);
    }
    request->isRecv = true;
    request->key    = key;
    node->addRequest(request);
    if (sock->asyncRecv(&request->buffer, request) || (lastErr = WSAGetLastError()) == WSA_IO_PENDING) {
        node->resetIdleCount();
        return true;
    }
    if (KSocket::isUdpSocket(*sock)) {
        KDebugLib("WSARecvFrom() failed: %d", lastErr);
    } else {
        KDebugLib("WSARecv() failed: %d", lastErr);
    }
    onError(sock, KDef::ErrorCallRecv, lastErr);
    node->removeRequest(request);
    request->key = NULL;
    m_pool.pushFront(request);
    delete key;
    return false;
}

bool KMultiplexor::disconnect(KSocket *sock)
{
    assert(sock != NULL && !KSocket::isUdpSocket(*sock));
    DWORD lastErr;
    KAutoPtr<KDef::DisconnRequest> request(new KDef::DisconnRequest);
    KDef::SocketContext *node = sock->getContext();
    node->addRequest(request);
    if (sock->asyncDisconnect(request) || (lastErr = WSAGetLastError()) == WSA_IO_PENDING) {
        request.release();
        node->resetIdleCount();
        return true;
    }
    KDebugLib("DisconnectEx() failed: %d", lastErr);
    onError(sock, KDef::ErrorCallDisconnect, lastErr);
    node->removeRequest(request);
    return false;
}

bool KMultiplexor::send(KSocket *sock, KStorage *storage, size_t length)
{
    assert(sock != NULL && storage != NULL && storage->bytesAvailable() > 0);
    assert(!KSocket::isUdpSocket(*sock));
    DWORD lastErr;
    KDef::IoRequest *request = static_cast<KDef::IoRequest *>(m_pool.popFront());
    KDef::SocketContext *node = sock->getContext();
    if (length > 0) {
        length = std::min<size_t>(storage->bytesAvailable(), length);
    } else {
        length = std::min<size_t>(storage->bytesAvailable(), m_chunkSize);
    }
    {
        size_t len = std::max<size_t>(length, m_chunkSize);
        if (request) {
            request->reset();
            request->resize(len);
        } else {
            request = new KDef::IoRequest(len);
            assert(request != NULL);
        }
    }
    request->bytesSent  = 0;
    request->bytesTotal = length;
    request->isRecv     = false;
    WSABUF buffer;
    buffer.buf = request->buffer.buf;
    buffer.len = length;
    storage->read(buffer.buf, length);
    node->addRequest(request);
    if (sock->asyncSend(&buffer, request) || (lastErr = WSAGetLastError()) == WSA_IO_PENDING) {
        node->resetIdleCount();
        return true;
    }
    KDebugLib("WSASend() failed: %d", lastErr);
    onError(sock, KDef::ErrorCallSend, lastErr);
    node->removeRequest(request);
    m_pool.pushFront(request);
    return false;
}

bool KMultiplexor::send(KSocket *sock, const void *data, size_t length, KKey *key)
{
    assert(sock != NULL && data != NULL && length > 0);
    assert(!KSocket::isUdpSocket(*sock));
    DWORD lastErr;
    KDef::IoRequest *request = static_cast<KDef::IoRequest *>(m_pool.popFront());
    KDef::SocketContext *node = sock->getContext();
    {
        size_t len = std::max<size_t>(length, m_chunkSize);
        if (request) {
            request->reset();
            request->resize(len);
        } else {
            request = new KDef::IoRequest(len);
            assert(request != NULL);
        }
    }
    request->bytesSent  = 0;
    request->bytesTotal = length;
    request->isRecv     = false;
    request->key        = key;
    WSABUF buffer;
    buffer.buf = request->buffer.buf;
    buffer.len = length;
    memmove(buffer.buf, data, length);
    node->addRequest(request);
    if (sock->asyncSend(&buffer, request) || (lastErr = WSAGetLastError()) == WSA_IO_PENDING) {
        node->resetIdleCount();
        return true;
    }
    KDebugLib("WSASend() failed: %d", lastErr);
    onError(sock, KDef::ErrorCallSend, lastErr);
    node->removeRequest(request);
    request->key = NULL;
    m_pool.pushFront(request);
    delete key;
    return false;
}

bool KMultiplexor::transmit(KSocket *sock, HANDLE hFile, size_t length, KKey *key, size_t blockLen, KPackage *head, KPackage *tail)
{
    assert(sock != NULL && hFile != INVALID_HANDLE_VALUE);
    assert(!KSocket::isUdpSocket(*sock));
    if (!sock->initTransmitFile()) {
        KDebugLib("Transmit failed to initialize");
        onError(sock, KDef::ErrorTransmitInit, WSAGetLastError());
        return false;
    }
    DWORD lastErr;
    KAutoPtr<KDef::TransmitRequest> request(new KDef::TransmitRequest(hFile, key));
    KDef::SocketContext *node = sock->getContext();
    node->addRequest(request);
    TRANSMIT_FILE_BUFFERS tfb = {0};
    if (head) {
        tfb.Head = head->data();
        tfb.HeadLength = head->length();
    }
    if (tail) {
        tfb.Tail = tail->data();
        tfb.TailLength = tail->length();
    }
    if (sock->asyncTransmit(hFile, length, blockLen, &tfb, request) || (lastErr = WSAGetLastError()) == WSA_IO_PENDING) {
        request.release();
        node->resetIdleCount();
        return true;
    }
    KDebugLib("TransmitFile() failed: %d", lastErr);
    onError(sock, KDef::ErrorCallTransmit, lastErr);
    node->removeRequest(request);
    return false;
}

size_t KMultiplexor::disconnectIf(const Condition &con)
{
    assert(con);
    KDef::SocketContext *context;
    KSocket *sock;
    size_t count = 0;
    {
        KLocker lock(m_context);
        KList::NodePtr node = m_context.front();
        while (node) {
            context = static_cast<KDef::SocketContext *>(node);
            sock = context->getSocket();
            if (KSocket::isUnspecific(*sock) && con(sock)) {
                if (disconnect(sock)) {
                    ++count;
                }
            }
            node = node->next();
        }
    }
    return count;
}

size_t KMultiplexor::sendIf(const void *data, size_t length, const Condition &con)
{
    assert(data != NULL && length > 0 && con);
    KDef::SocketContext *context;
    KSocket *sock;
    size_t count = 0;
    {
        KLocker lock(m_context);
        KList::NodePtr node = m_context.front();
        while (node) {
            context = static_cast<KDef::SocketContext *>(node);
            sock = context->getSocket();
            if (KSocket::isUnspecific(*sock) && con(sock)) {
                if (send(sock, data, length)) {
                    ++count;
                }
            }
            node = node->next();
        }
    }
    return count;
}

bool KMultiplexor::sendTo(KSocket *sock, const KIAddress &addrRemote, const void *data, size_t length, KKey *key)
{
    assert(sock != NULL && !addrRemote.isNull() && data != NULL && length > 0);
    assert(KSocket::isUdpSocket(*sock));
    DWORD lastErr;
    KDef::IoRequest *request = static_cast<KDef::IoRequest *>(m_pool.popFront());
    KUdpSocketContext *node = static_cast<KUdpSocketContext *>(sock);
    {
        size_t len = std::max<size_t>(length, m_chunkSize);
        if (request) {
            request->reset();
            request->resize(len);
        } else {
            request = new KDef::IoRequest(len);
            assert(request != NULL);
        }
    }
    request->addrTo     = *reinterpret_cast<const SOCKADDR_IN *>(addrRemote.name());
    request->bytesSent  = 0;
    request->bytesTotal = length;
    request->isRecv     = false;
    request->key        = key;
    WSABUF buffer;
    buffer.buf = request->buffer.buf;
    buffer.len = length;
    memmove(buffer.buf, data, length);
    node->addRequest(request);
    if (sock->asyncSend(&buffer, request) || (lastErr = WSAGetLastError()) == WSA_IO_PENDING) {
        node->resetIdleCount();
        return true;
    }
    KDebugLib("WSASendTo() failed: %d", lastErr);
    onError(sock, KDef::ErrorCallSend, lastErr);
    node->removeRequest(request);
    request->key = NULL;
    m_pool.pushFront(request);
    delete key;
    return false;
}

void KMultiplexor::ensureAlive(KSocket *sock, bool always)
{
    deleteLater(sock);
    if (always) {
        sock->setRequestTimeOut(INFINITE);
    } else if (sock->requestTimeOut() == INFINITE) {
        sock->setRequestTimeOut();
    }
}

KDef::Life KMultiplexor::lifeBegin(KSocket *sock)
{
    assert(sock != NULL);
    KDef::NullRequest *request = new KDef::NullRequest;
    KDef::SocketContext *node = sock->getContext();
    assert(request != NULL);
    node->addRequest(request);
    node->resetIdleCount();
    return request;
}

void KMultiplexor::lifeEnd(KSocket *sock, KDef::Life life)
{
    assert(sock != NULL && life != NULL);
    KDef::NullRequest *request = static_cast<KDef::NullRequest *>(life);
    KDef::SocketContext *node = sock->getContext();
    node->removeRequest(request);
    delete request;
}

bool KMultiplexor::startImp()
{
    if (m_stopped) {
        m_iocp.reset(new KSocketIocp);
        for (size_t i = 0; i < m_nThreads; ++i) {
            m_workers[i].start();
        }
        m_gc->start();
        m_stopped = false;
        return true;
    }
    return false;
}

bool KMultiplexor::stopImp(DWORD ms)
{
    if (!m_stopped) {
        onStopping();
        if (m_gc->isRunning()) {
            m_gcStop.set();
            m_gc->join();
        }
        for (size_t i = 0; i < m_nThreads; ++i) {
            if (!m_iocp->postStatus(KSocketIocp::Status(0, 0, NULL))) {
                onError(NULL, KDef::ErrorPostStatus, GetLastError());
                return false;
            }
        }
        if (!KThread::waitAll(m_workers.get(), m_nThreads, ms)) {
            KDebugLib("Workers failed to stop in time");
            onError(NULL, KDef::ErrorWaitTimeOut, GetLastError());
            return false;
        }
        {
            KLocker lock(m_context);
            KList::NodePtr node = m_context.front();
            while (node) {
                onClosing(static_cast<KDef::SocketContext *>(node)->getSocket(), true);
                node = node->next();
            }
        }
        m_context.destroy();
        m_stopped = true;
        return true;
    }
    return false;
}

void KMultiplexor::_gc()
{
    while (!m_gcStop.wait(KDef::DefaultInterval)) {
        KLocker lock(m_context);
        KList::NodePtr node = m_context.front();
        while (node) {
            node = validate(static_cast<KDef::SocketContext *>(node));
        }
    }
}

void KMultiplexor::_worker()
{
    KDef::RequestInfo info(this);
    LPWSAOVERLAPPED request;
    for (;;) {
        info.node = NULL;
        request = NULL;
        if (m_iocp->getStatus(&info.bytes, reinterpret_cast<PULONG_PTR>(&info.node), &request)) {
            if (info.node == NULL) {
                break;
            }
            static_cast<KDef::BasicRequest *>(request)->onSuccess(info);
        } else if ((info.error = GetLastError()) == ERROR_MORE_DATA) {
            KDebugLib("More data is available...");
            assert(request != NULL && info.node != NULL);
            info.more = true;
            static_cast<KDef::BasicRequest *>(request)->onSuccess(info);
            info.more = false;
        } else {
            KDebugLib("GetQueuedCompletionStatus() failed: %d", info.error);
            if (request) {
                assert(info.node != NULL);
                static_cast<KDef::BasicRequest *>(request)->onFailure(info);
            } else {
                onError(NULL, KDef::ErrorGetStatus, info.error);
            }
        }
    }
    KDebugLib("Worker quitted");
}

bool KMultiplexor::accept(KSocket *sock, KDef::AcceptRequest *request)
{
    KTcpSocketContext *node = static_cast<KTcpSocketContext *>(sock);
    if (request) {
        request->reset();
    } else {
        request = new KDef::AcceptRequest;
    }
    KAutoPtr<KDef::AcceptRequest> holder(request);
    KAutoPtr<KSocketContext> client(new KSocketContext);
    if (client->open() && client->initDisconnectEx()) {
        DWORD lastErr;
        request->client = client.release();
        node->addRequest(request);
        if (sock->asyncAccept(request->client, request->buf, request) || (lastErr = WSAGetLastError()) == WSA_IO_PENDING) {
            holder.release();
            node->resetIdleCount();
            return true;
        }
        KDebugLib("AcceptEx() failed: %d", lastErr);
        onError(sock, KDef::ErrorCallAccept, lastErr);
        node->removeRequest(request);
        return false;
    }
    KDebugLib("Client failed to initialize");
    onError(sock, KDef::ErrorClientInit, WSAGetLastError());
    return false;
}

bool KMultiplexor::connect(KSocket *sock, const KIAddress &addr)
{
    DWORD lastErr;
    KAutoPtr<KDef::ConnectRequest> request(new KDef::ConnectRequest);
    KTcpSocketContext *node = static_cast<KTcpSocketContext *>(sock);
    node->addRequest(request);
    if (sock->asyncConnect(addr, request) || (lastErr = WSAGetLastError()) == WSA_IO_PENDING) {
        request.release();
        node->resetIdleCount();
        return true;
    }
    KDebugLib("ConnectEx() failed: %d", lastErr);
    onError(sock, KDef::ErrorCallConnect, lastErr);
    node->removeRequest(request);
    return false;
}

void KMultiplexor::deleteLater(KSocket *sock)
{
    assert(sock != NULL);
    KDef::SocketContext *node = sock->getContext();
    node->resetIdleCount();
}

bool KMultiplexor::sendAgain(KSocket *sock, KDef::BasicRequest *request, WSABUF *buffer)
{
    KDebugLib("Warning: KMultiplexor::sendAgain() called");
    DWORD lastErr;
    if (sock->asyncSend(buffer, request) || (lastErr = WSAGetLastError()) == WSA_IO_PENDING) {
        KDef::SocketContext *node = sock->getContext();
        node->resetIdleCount();
        return true;
    }
    if (KSocket::isUdpSocket(*sock)) {
        KDebugLib("WSASendTo() failed: %d", lastErr);
    } else {
        KDebugLib("WSASend() failed: %d", lastErr);
    }
    onError(sock, KDef::ErrorCallSend, lastErr);
    return false;
}

KList::NodePtr KMultiplexor::validate(KDef::SocketContext *node)
{
    KList::NodePtr next = node->next();
    KSocket *sock = node->getSocket();
    DWORD timeOut = sock->requestTimeOut();
    if (timeOut < INFINITE) {
        if (node->hasPendingRequest()) {
            if (KSocket::s_cancel && timeOut > 0) {
                if (DWORD(node->increaseIdleCount() * KDef::DefaultInterval) >= timeOut) {
                    cancel(sock);
                }
            } else if (node->isIdle()) {
                node->resetIdleCount();
            }
        } else if (node->increaseIdleCount() >= KDef::DefaultExpireCount) {
            m_context.cutOffUnlocked(node);
            onClosing(node->getSocket(), false);
            delete node;
        }
    }
    return next;
}